package com.bupt.ilink.tool;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.Constants;

import java.util.ArrayList;
import java.util.List;

public class BailianModelTools {
    public static GenerationParam createGenerationParam(List<Message> messages) {
        return GenerationParam.builder()
                .model("qwen2-72b-instruct")
                .messages(messages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8)
                .build();
    }

    public static GenerationResult callGenerationWithMessages(GenerationParam param) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        return gen.call(param);
    }

    public static String runLLM(String prompt, String userInput) {
        try {
            List<Message> messages = new ArrayList<>();
            Constants.apiKey="sk-36ef0e3632a545c38af1c48164b19769";

            System.out.println("模型输入："+prompt);
            System.out.println("用户输入："+userInput);

            messages.add(createMessage(Role.SYSTEM, prompt));
            messages.add(createMessage(Role.USER, userInput));
            GenerationParam param = createGenerationParam(messages);
            GenerationResult result = callGenerationWithMessages(param);
            System.out.println("模型输出："+result.getOutput().getChoices().get(0).getMessage().getContent());

//            messages.add(result.getOutput().getChoices().get(0).getMessage());
            return result.getOutput().getChoices().get(0).getMessage().getContent();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            e.printStackTrace();
            return null;
        }
    }

    public static String buildPrompt(String context) {
        String prompt = "请根据下列```内的内容回答问题，并确保回答为json格式。\n"
                + "请将回答组织成包含'答案'和'参考来源'两个字段的json格式，具体格式如下：\n"
                + "{\n"
                + "  \"答案\": \"[在此填写具体回答]\",\n"
                + "  \"参考来源\": \"[在此填写相关参考来源，确保与检索到的'参考来源'字段内容一致]\"\n"
                + "}\n"
                + "内容如下：\n"
                + "```\n"
                + context + "\n"
                + "```\n"
                + "请确保回答的结构和内容完整，格式正确。"
                + "注意若以上提示对回答问题没有帮助，则自己组织答案，并在‘参考来源’处填写‘自行组织’。";
        return prompt;
    }

    public static String buildUserInput(String query) {
        String userInput = "我的问题是:" + query;
        return userInput;
    }


    private static Message createMessage(Role role, String content) {
        return Message.builder().role(role.getValue()).content(content).build();
    }
}
